#include "wmi_registrymanager.h"
#include "wmi_hresult.h"

using namespace log4cxx;
using namespace log4cxx::xml;
using namespace log4cxx::helpers;
using namespace std;
WMI_RegistryManager::WMI_RegistryManager(WMI_Connection* con)
{
    DOMConfigurator::configure("LogConfig.xml");


    LoggerPtr loggerRegMan(Logger::getLogger("loggerRegMan"));
    stringstream out;

    _con = con;
    _ClassName = SysAllocString(STDREGPROV.c_str());
    WMI_HResult* whres = _con->getClassObject(_ClassName, &_pClass);

    if (FAILED(whres->getHRESULT()))
    {
        out << NOTLOADCLASS.c_str() << _ClassName << " Error code = 0x" 
            << hex << whres->getHRESULT() << endl;
        LOG4CXX_ERROR(loggerRegMan,out.str());
        _pClass->Release();
        CoUninitialize();
    }

    delete whres;
}

WMI_RegistryManager::~WMI_RegistryManager(void)
{

}

WMI_HResult*
WMI_RegistryManager::getRegistrys(WMI_Registry** rClass)
{
    wstring root;
    wstring subKey;

    getParseRoot((*rClass)->getPath(), root, subKey);
    return getRegistrys(convertStringToHKEY(root), subKey, rClass);
}

WMI_HResult* 
WMI_RegistryManager::getRegistrys(HiveEnum root, wstring SubKey, WMI_Registry** rClass)
{
    HRESULT hres;    
    LoggerPtr loggerRegMan(Logger::getLogger("loggerRegMan"));
    stringstream out;

    IWbemClassObject* pOutParams = NULL;
    getMethod(root, SubKey, ENUMKEY, &pOutParams,VALUE_NAME);
    LoadValues(rClass);

    VARIANT varReturnValue;
    varReturnValue.vt = VT_BSTR;
    hres = pOutParams->Get(_bstr_t(SNAMES.c_str()), 0,&varReturnValue, NULL, 0);

    if(varReturnValue.vt == 1)
    {
        WMI_HResult* whresult = new WMI_HResult(NOTHAVECHILDREN, -1);
        LOG4CXX_ERROR(loggerRegMan, whresult->toString());
        return whresult;
    }

    LONG lstart, lend;
    LONG idx = -1;
    LONG nPos;	
    BSTR* pbstr;

    // assign the Safearray
    SAFEARRAY *sa = V_ARRAY(&varReturnValue); 

    // Get the lower and upper bound
    hres = SafeArrayGetLBound(sa, 1, &lstart);

    if(FAILED(hres))
    {
        return new WMI_HResult(L"", hres);
    }

    hres = SafeArrayGetUBound( sa, 1, &lend );

    if(FAILED(hres))
    {
        return new WMI_HResult(L"", hres);
    }
    // loop
    hres = SafeArrayAccessData(sa, (void HUGEP**)&pbstr);

    if(SUCCEEDED(hres))
    {
        for(idx = lstart; idx <= lend; idx++)
        {		
            BSTR s;
            s = pbstr[idx];
            wstring name = BSTR(s);
            WMI_Registry* auxclass;

            if(SubKey.empty())
            {				
                auxclass = new WMI_Registry(name, convertHiveEnumToWstring(root));				
            }
            else
            {
                wstring sub = convertHiveEnumToWstring(root);
                sub += WSBACKSLASH;
                sub += SubKey;
                sub += WSBACKSLASH;
                sub += name;				
                auxclass = new WMI_Registry(name, sub);				
            }

            (*rClass)->addRegistry(auxclass);
        }

        hres = SafeArrayUnaccessData(sa);	

        if(FAILED(hres))
        {
            return new WMI_HResult(L"", hres);
        }
    }

    return new WMI_HResult(W_OK, 0);;
}

WMI_HResult* 
WMI_RegistryManager::LoadValues(WMI_Registry** rClass)
{	
    HRESULT hres;
    WMI_HResult* whres;
    LoggerPtr loggerRegMan(Logger::getLogger("loggerRegMan"));
    stringstream out;

    IWbemClassObject* pOutParams = NULL;

    wstring root;
    wstring subKey;

    getParseRoot((*rClass)->getPath(), root, subKey);

    getMethod(convertStringToHKEY(root), subKey, ENUMVALUES, &pOutParams, VALUE_NAME);

    CIMTYPE pvtType1 = NULL;
    CIMTYPE pvtType2 = NULL;

    VARIANT varReturnValue;
    varReturnValue.vt = VT_DISPATCH;
    hres = pOutParams->Get(_bstr_t(SNAMES.c_str()), 0, &varReturnValue, &pvtType1, 0);

    if (FAILED(hres))
    {
        whres = new WMI_HResult(NOTGETSNAMES, hres);
        LOG4CXX_ERROR(loggerRegMan, whres->toString());
        return whres;
    }

    if(varReturnValue.vt == 1)
    {	
        return new WMI_HResult(L"", -1);
    }

    LONG lstart, lend;
    LONG idx = -1;
    LONG nPos;	
    BSTR* pbstr;

    // assign the Safearray
    SAFEARRAY *sa = V_ARRAY(&varReturnValue); 

    // Get the lower and upper bound
    hres = SafeArrayGetLBound(sa, 1, &lstart);

    if(FAILED(hres))
    {
        whres = new WMI_HResult(NOTGETLBOUNDARRAY, hres);
        LOG4CXX_ERROR(loggerRegMan, whres->toString());
        return whres;
    }

    hres = SafeArrayGetUBound(sa, 1, &lend);

    if(FAILED(hres))
    {
        whres = new WMI_HResult(NOTGETUBOUNDARRAY, hres);
        LOG4CXX_ERROR(loggerRegMan, whres->toString());
        return whres;
    }

    // loop
    hres = SafeArrayAccessData(sa,(void HUGEP**)&pbstr);

    if(SUCCEEDED(hres))
    {
        for(idx = lstart; idx <= lend; idx++)
        {		
            BSTR s;
            s = pbstr[idx];
            wstring name = BSTR(s);
            WMI_Value* auxclass;

            auxclass = new WMI_Value(name, EMPTY_TYPE);		
            (*rClass)->addValue(*auxclass);

            delete auxclass;
        }

        hres = SafeArrayUnaccessData(sa);	

        if(FAILED(hres))
        {
            return new WMI_HResult(L"", hres);
        }
    }	

    return new WMI_HResult(W_OK, 0);
}

WMI_HResult*
WMI_RegistryManager::getMethod(HiveEnum root, wstring subKey, wstring nameMethod, IWbemClassObject** pOutParams, wstring ValueName)
{
    HRESULT hres;
    WMI_HResult* whres;
    LoggerPtr loggerRegMan(Logger::getLogger("loggerRegMan"));
    stringstream out;

    BSTR MethodName = SysAllocString(nameMethod.c_str());

    IWbemClassObject* pInParamsDefinition = NULL;
    hres = _pClass->GetMethod(MethodName, 0, &pInParamsDefinition, NULL);

    if (FAILED(hres))
    {
        whres = new WMI_HResult(NOTRETURNMETHOD, hres);
        //out << NOTRETURNMETHOD.c_str() << nameMethod.c_str() << "' Error code = 0x" << hex << hres << endl;
        LOG4CXX_ERROR(loggerRegMan, whres->toString());
        return whres;
    }

    IWbemClassObject* pClassInstance = NULL;
    hres = pInParamsDefinition->SpawnInstance(0, &pClassInstance);


    if (FAILED(hres))
    {
        //WMI_HResult* whres = new WMI_HResult(L"", NOTCREATEINSTANCE, hres);
        out.str("");
        out << NOTCREATEINSTANCE.c_str() << "' Error code = 0x" << hex << hres << endl;
        LOG4CXX_ERROR(loggerRegMan,out.str());
    }

    // Create the values for the in parameters
    VARIANT varDefKey;
    varDefKey.vt = VT_I4;

    switch (root)
    {	
    case R_CLASSES_ROOT:
        varDefKey.lVal = (long) HKEY_CLASSES_ROOT;
        break;

    case R_CURRENT_USER:
        varDefKey.lVal = (long) HKEY_CURRENT_USER;
        break;

    case R_LOCAL_MACHINE:
        varDefKey.lVal = (long) HKEY_LOCAL_MACHINE;
        break;

    case R_USERS:
        varDefKey.lVal = (long) HKEY_USERS;
        break;

    case R_CURRENT_CONFIG:
        varDefKey.lVal = (long) HKEY_CURRENT_CONFIG;
        break;

    case R_DYN_DATA:
        varDefKey.lVal = (long) HKEY_DYN_DATA;
        break;
    }

    // Store the value for the in parameters
    hres = pClassInstance->Put(HDEFKEY.c_str(), 0,&varDefKey, 0);

    if (FAILED(hres))
    {
        whres = new WMI_HResult(NOTPUTHDEFKEY, hres);
        LOG4CXX_ERROR(loggerRegMan, whres->getHRESULT());
        return whres;
    }

    // Create the values for the in parameters
    VARIANT varSubKeyName;
    varSubKeyName.vt = VT_BSTR;
    varSubKeyName.bstrVal = (BSTR) subKey.c_str();

    hres = pClassInstance->Put(SSUBKEYNAME.c_str(), 0, &varSubKeyName, 0);

    if (FAILED(hres))
    {
        whres = new WMI_HResult(NOTPUTSSUBKEYNAME, hres);
        LOG4CXX_ERROR(loggerRegMan, whres->toString());
        return whres;
    }

    if(!ValueName.empty())
    {
        VARIANT varValueName;
        varValueName.vt = VT_BSTR;
        varValueName.bstrVal = (BSTR) ValueName.c_str();

        hres = pClassInstance->Put(SVALUENAME.c_str(), 0, &varValueName, 0);

        if (FAILED(hres))
        {
            whres = new WMI_HResult(NOTPUTSVALUENAME, hres);
            LOG4CXX_ERROR(loggerRegMan, whres->toString());
            return whres;
        }
    }

    hres = _con->getServices()->ExecMethod(_ClassName, MethodName,0,
        NULL, pClassInstance, pOutParams, NULL);

    if (FAILED(hres))
    {
        whres = new WMI_HResult(NOTEXECUTEMETHOD, hres);
        LOG4CXX_ERROR(loggerRegMan, whres->toString());
        return whres;
    }

    return new WMI_HResult(W_OK, 0);
}

WMI_HResult*
WMI_RegistryManager::getStringValue(HiveEnum root, wstring subKey, wstring valueName, wstring valueString)
{
    HRESULT hres;
    WMI_HResult* whres;
    LoggerPtr loggerRegMan(Logger::getLogger("loggerRegMan"));
    stringstream out;

    IWbemClassObject* pOutParams = NULL;
    getMethod(root, subKey, GETSTRINGVALUE, &pOutParams, valueName);

    VARIANT varReturnValue;
    varReturnValue.vt = VT_BSTR;

    hres = pOutParams->Get(_bstr_t(SVALUE.c_str()), 0, &varReturnValue, NULL, 0);

    if (FAILED(hres))
    {
        whres = new WMI_HResult(NOTGETSVALUE, hres);
        LOG4CXX_ERROR(loggerRegMan, whres->toString());
        return whres;
    }

    valueString = varReturnValue.bstrVal;
    return new WMI_HResult(W_OK, 0);
}

int 
WMI_RegistryManager::getParseRoot(std::wstring path, std::wstring &root, std::wstring &subKey)
{
    int cont = 0;
    for(int i = 0; i < path.size(); i++)
    {
        if(path[i] != BACKSLASH)
        {
            cont++;
        }
        else
        {
            break;
        }
    }

    wstring aux = path.substr(0, cont);
    root = path.substr(0, cont);
    if (cont == path.size())
    {
        subKey = EMPTY_SUBKEY;
    }
    else
    {
        subKey = path.substr(cont + 1, path.size());
    }

    return cont;
}

HiveEnum
WMI_RegistryManager::convertStringToHKEY(wstring strKey)
{
    if(strKey.compare(HKCR) || strKey.compare(CONST_HKEY_CLASSES_ROOT))
        return R_CLASSES_ROOT;

    else if(strKey.compare(HKCU) || strKey.compare(CONST_HKEY_CURRENT_USER))
        return R_CURRENT_USER;

    else if(strKey.compare(HKLM) || strKey.compare(CONST_HKEY_LOCAL))
        return R_LOCAL_MACHINE;

    else if(strKey.compare(HKU) || strKey.compare(CONST_HKEY_USERS))			
        return R_USERS;

    else if(strKey.compare(HKCF) || strKey.compare(CONST_HKEY_CURRENT_CONFIG))			
        return R_CURRENT_CONFIG;

    else if(strKey.compare(HKDD) || strKey.compare(CONST_HKEY_DYN_DATA))			
        return R_DYN_DATA;
}

wstring 
WMI_RegistryManager::convertHiveEnumToWstring(HiveEnum hive)
{
    wstring result;
    switch (hive)
    {	
    case R_CLASSES_ROOT:
        result = CONST_HKEY_CLASSES_ROOT;
        break;

    case R_CURRENT_USER:
        result = CONST_HKEY_CURRENT_USER;
        break;

    case R_LOCAL_MACHINE:
        result = CONST_HKEY_LOCAL;
        break;

    case R_USERS:
        result = CONST_HKEY_USERS;
        break;

    case R_CURRENT_CONFIG:
        result = CONST_HKEY_CURRENT_CONFIG;
        break;

    case R_DYN_DATA:
        result = CONST_HKEY_DYN_DATA;
        break;
    }

    return result;
}

void
WMI_RegistryManager::getStandardRegistryHives(std::vector<WMI_Registry>& standarRegistrys)
{

    WMI_Registry reg(CONST_HKEY_CLASSES_ROOT,convertHiveEnumToWstring(R_CLASSES_ROOT));
    standarRegistrys.push_back(reg);

    reg = WMI_Registry(CONST_HKEY_CURRENT_CONFIG,convertHiveEnumToWstring(R_CURRENT_CONFIG));
    standarRegistrys.push_back(reg);

    reg = WMI_Registry(CONST_HKEY_CURRENT_USER,convertHiveEnumToWstring(R_CURRENT_USER));
    standarRegistrys.push_back(reg);

    reg = WMI_Registry(CONST_HKEY_LOCAL,convertHiveEnumToWstring(R_LOCAL_MACHINE));
    standarRegistrys.push_back(reg);

}

void
WMI_RegistryManager::getRegistryName(const wstring& path, wstring& outRoot, wstring& outRegistryName)
{
    int name = path.rfind(L"\\", path.size() - 1);
        if(name > 0)
    {
	    outRoot = path.substr(0, name);
    }
    else
    {
        outRoot.clear();
    }
	outRegistryName = path.substr(name + 1);
	
}

void 
WMI_RegistryManager::getRegistrySubkey(const wstring& path, wstring& outRegistryName)
{
    int name = path.rfind(L"\\", path.size() - 1);
    if(name > 0)
    {
	    outRegistryName = path.substr(name + 1);
    }
    else
    {
        outRegistryName.clear();
    }
}